import os
from typing import Optional, List

from fit_py_ask_number_demo.data_managers import sqlite_manager
from fit_py_ask_number_demo.configs import config


class PromptManager:
    def __init__(self):
        """
        初始化 Prompt 管理器
        """
        db_path = os.path.join(os.path.dirname(__file__), '../../database/kpi2.db')
        self.db_manager = sqlite_manager.SQLiteManager(db_path)

    def initialize_table(self):
        """
        初始化 prompt 表，并插入三行数据
        如果表已存在，则跳过初始化
        """
        if not self.db_manager.connection:
            self.db_manager.connect()

        try:
            cursor = self.db_manager.connection.cursor()

            # 检查表是否存在
            cursor.execute('''
            SELECT name FROM sqlite_master WHERE type='table' AND name='prompt'
            ''')
            table_exists = cursor.fetchone()

            if table_exists:
                return

            # 创建 prompt 表
            cursor.execute('''
            CREATE TABLE prompt (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT NOT NULL,
                prompt TEXT NOT NULL
            )
            ''')

            # 插入三行数据
            initial_data = [
                {'name': '问题路由', 'prompt': config.question_route_prompt},
                {'name': '信息提取', 'prompt': config.info_extract_prompt},
                {'name': 'DSL生成', 'prompt': config.dsl_generate_prompt}
            ]
            for data in initial_data:
                self.db_manager.create('prompt', data)

            self.db_manager.connection.commit()
            print("prompt 表初始化成功！")
        except Exception as e:
            print(f"初始化 prompt 表失败: {e}")

    def get_prompt_by_name(self, name: str) -> Optional[str]:
        """
        根据 name 查询 prompt 值
        :param name: 名称
        :return: prompt 值，如果未找到则返回 None
        """
        result = self.db_manager.read('prompt', f"name = '{name}'")
        if result:
            return result[0]['prompt']
        return None

    def update_prompt_by_name(self, name: str, new_prompt: str) -> bool:
        """
        根据 name 更新 prompt 值
        :param name: 名称
        :param new_prompt: 新的 prompt 值
        :return: 是否成功
        """
        return self.db_manager.update('prompt', {'prompt': new_prompt}, f"name = '{name}'")

    def list_all_prompts(self) -> list[dict]:
        """
        查看表中所有的 name 和对应的 prompt 值
        :return: 包含所有记录的列表，每行数据为一个字典
        """
        prompt_list = []
        all_prompts = self.db_manager.read('prompt')
        for prompt in all_prompts:
            prompt_list.append({"name": prompt["name"], "prompt": prompt["prompt"]})
        return prompt_list


prompt_manager = PromptManager()
